#!/bin/bash

set -e

# Default model configuration
DEFAULT_MODEL_ID="LiquidAI/LFM2-1.2B"

CLI_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_ROOT="$(dirname "$CLI_DIR")"
BUILD_DIR="$CLI_DIR/build"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Check for required dependencies
check_dependencies() {
    local missing_packages=()

    # Check for CMake
    if ! command -v cmake &> /dev/null; then
        echo -e "${YELLOW}CMake is not installed${NC}"
        missing_packages+=("cmake")
    fi

    # Check for Python3
    if ! command -v python3 &> /dev/null; then
        echo -e "${YELLOW}Python3 is not installed${NC}"
        missing_packages+=("python3")
    fi

    # If packages are missing, handle installation
    if [ ${#missing_packages[@]} -gt 0 ]; then
        echo -e "\n${YELLOW}Missing packages: ${missing_packages[*]}${NC}"

        # Check if brew is available
        if command -v brew &> /dev/null; then
            echo -e "\n${BLUE}Homebrew is available.${NC}"
            echo -n "Would you like to install the missing packages with brew? [y/N]: "
            read -r response

            if [[ "$response" =~ ^[Yy]$ ]]; then
                for package in "${missing_packages[@]}"; do
                    echo -e "\n${BLUE}Installing $package...${NC}"
                    brew install "$package"
                done
                echo -e "\n${GREEN}Installation complete!${NC}"
            else
                echo -e "\n${YELLOW}Please install the missing packages manually:${NC}"
                for package in "${missing_packages[@]}"; do
                    echo "  brew install $package"
                done
                exit 1
            fi
        else
            echo -e "\n${YELLOW}Homebrew is not installed.${NC}"
            echo "To install the missing packages, you can either:"
            echo ""
            echo "1. Install Homebrew first:"
            echo '   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"'
            echo "   Then run: brew install ${missing_packages[*]}"
            echo ""
            echo "2. Or install the packages manually using your system's package manager"
            exit 1
        fi
    else
        # All packages exist, now check for Python venv
        setup_python_venv
    fi
}

# Setup Python virtual environment and install dependencies
setup_python_venv() {
    VENV_DIR="$PROJECT_ROOT/venv"
    REQUIREMENTS_FILE="$PROJECT_ROOT/tools/requirements.txt"

    # Check if venv exists
    if [ ! -d "$VENV_DIR" ]; then
        echo -e "${BLUE}Setting up Python virtual environment...${NC}"

        # Create venv
        if ! python3 -m venv "$VENV_DIR"; then
            echo -e "${RED}Failed to create virtual environment${NC}"
            echo "Please ensure python3-venv is installed:"
            echo "  Ubuntu/Debian: sudo apt-get install python3-venv"
            echo "  macOS: Should be included with Python3"
            exit 1
        fi

        echo -e "${GREEN}Virtual environment created${NC}"
    fi

    # Activate venv
    source "$VENV_DIR/bin/activate"

    # Check if requirements file exists
    if [ -f "$REQUIREMENTS_FILE" ]; then
        # Check if dependencies are already installed
        if ! python -c "import numpy, torch, transformers" 2>/dev/null; then
            echo -e "${BLUE}Installing Python dependencies...${NC}"
            echo "This may take a few minutes on first run."

            # Upgrade pip first
            pip install --upgrade pip >/dev/null 2>&1

            # Install requirements
            if pip install -r "$REQUIREMENTS_FILE"; then
                echo -e "${GREEN}Python dependencies installed successfully${NC}"
            else
                echo -e "${RED}Failed to install Python dependencies${NC}"
                exit 1
            fi
        else
            echo -e "${GREEN}Python dependencies already installed${NC}"
        fi
    else
        echo -e "${YELLOW}Warning: requirements.txt not found at $REQUIREMENTS_FILE${NC}"
        echo "Python packages will need to be installed manually if needed"
    fi
}

print_usage() {
    echo "Usage: $0 <command> [options]"
    echo ""
    echo "Commands:"
    echo "  build         - Build the chat application only"
    echo "  download      - Download model weights only"
    echo "  run           - Build, download (if needed), and run the chat app"
    echo "  test          - Run the test suite"
    echo ""
    echo "Arguments:"
    echo "  model_id      - HuggingFace model ID (e.g., LiquidAI/LFM2-1.2B)"
    echo ""
    echo "Test options:"
    echo "  --model <name>            Model to use for tests (default: Qwen/Qwen3-0.6B)"
    echo "  --transcribe_model <name> Transcribe model to use (default: openai/whisper-small)"
    echo ""
    echo "Examples:"
    echo "  $0 build"
    echo "  $0 download google/gemma-3-270m-it"
    echo "  $0 run"
    echo "  $0 run Qwen/Qwen3-0.6B"
    echo "  $0 test"
    echo "  $0 test --model Qwen/Qwen3-0.6B"
    echo "  $0 test --transcribe_model openai/whisper-large"
    echo "  $0 test --model LiquidAI/LFM2-VL-450M --transcribe_model openai/whisper-large"
    echo ""
    echo "Common models:"
    echo "  LiquidAI/LFM2-1.2B (default)"
    echo "  Qwen/Qwen3-0.6B"
    echo "  google/gemma-3-270m-it"
    echo "  google/gemma-3-1b-it"
}

get_model_dir_name() {
    # Convert model ID to directory name
    # e.g., "LiquidAI/LFM2-1.2B" -> "lfm2-1.2B"
    # e.g., "Qwen/Qwen3-0.6B" -> "qwen3-0.6B"
    # e.g., "google/gemma-3-270m-it" -> "gemma3-270m"

    local model_id="$1"
    local model_name=$(echo "$model_id" | cut -d'/' -f2)

    # Convert to lowercase and clean up
    model_name=$(echo "$model_name" | tr '[:upper:]' '[:lower:]' | sed 's/-it$//')

    echo "$model_name"
}

build_chat() {
    echo -e "${BLUE}Building Cactus chat...${NC}"
    echo "======================="

    mkdir -p "$BUILD_DIR"

    # Build Cactus library if needed
    cd "$PROJECT_ROOT/cactus"
    if [ ! -f "build/libcactus.a" ]; then
        echo -e "${YELLOW}Cactus library not found. Building it now...${NC}"
        if [ -f "$PROJECT_ROOT/cactus/build.sh" ]; then
            "$PROJECT_ROOT/cactus/build.sh"
        else
            echo -e "${RED}Error: build.sh not found at $PROJECT_ROOT/cactus/build.sh${NC}"
            exit 1
        fi
    else
        echo -e "${GREEN}Cactus library found.${NC}"
    fi

    cd "$BUILD_DIR"

    echo "Compiling chat.cpp..."

    if [[ "$OSTYPE" == "darwin"* ]]; then
        clang++ -std=c++17 -O3 \
            -I"$PROJECT_ROOT" \
            "$CLI_DIR/chat.cpp" \
            "$PROJECT_ROOT/cactus/build/libcactus.a" \
            -o chat \
            -framework Accelerate
    else
        g++ -std=c++17 -O3 \
            -I"$PROJECT_ROOT" \
            "$CLI_DIR/chat.cpp" \
            "$PROJECT_ROOT/cactus/build/libcactus.a" \
            -o chat \
            -pthread
    fi

    echo -e "${GREEN}Build complete: $BUILD_DIR/chat${NC}"
    echo ""
}

download_model() {
    local model_id="$1"
    local model_dir=$(get_model_dir_name "$model_id")
    local weights_dir="$PROJECT_ROOT/weights/$model_dir"

    if [ -d "$weights_dir" ] && [ -f "$weights_dir/config.txt" ]; then
        echo -e "${GREEN}Model weights found at $weights_dir${NC}"
        return 0
    fi

    echo ""
    echo -e "${YELLOW}Model weights not found. Downloading $model_id...${NC}"
    echo "============================================="

    cd "$PROJECT_ROOT"

    if ! command -v python3 &> /dev/null; then
        echo -e "${RED}Error: Python3 not found. Cannot download weights automatically.${NC}"
        echo "Please run manually: python3 tools/convert_hf.py $model_id weights/$model_dir/ --precision INT8"
        return 1
    fi

    if ! python3 -c "import numpy, torch, transformers" 2>/dev/null; then
        echo -e "${RED}Error: Required Python packages not found.${NC}"
        echo "Please check the README for setup instructions."
        return 1
    fi

    echo "Running: python3 tools/convert_hf.py $model_id weights/$model_dir/ --precision INT8"
    if python3 tools/convert_hf.py "$model_id" "weights/$model_dir/" --precision INT8; then
        echo -e "${GREEN}Successfully downloaded and converted weights${NC}"
        return 0
    else
        echo -e "${RED}Error: Failed to download weights.${NC}"
        echo "Please run manually: python3 tools/convert_hf.py $model_id weights/$model_dir/ --precision INT8"
        return 1
    fi
}

run_chat() {
    local model_id="${1:-$DEFAULT_MODEL_ID}"
    local model_dir=$(get_model_dir_name "$model_id")
    local weights_dir="$PROJECT_ROOT/weights/$model_dir"

    # Build the chat app
    build_chat

    # Download model if needed
    if ! download_model "$model_id"; then
        exit 1
    fi

    # Run the chat app
    clear
    echo -e "${GREEN}Starting Cactus Chat with model: $model_id${NC}"
    echo ""

    "$BUILD_DIR/chat" "$weights_dir"
}

run_tests() {
    local model_arg=""
    local transcribe_model_arg=""

    # Parse test-specific arguments
    shift  # Remove 'test' command
    while [[ $# -gt 0 ]]; do
        case $1 in
            --model)
                model_arg="--model $2"
                shift 2
                ;;
            --transcribe_model)
                transcribe_model_arg="--transcribe_model $2"
                shift 2
                ;;
            *)
                echo -e "${RED}Unknown test option: $1${NC}"
                exit 1
                ;;
        esac
    done

    echo -e "${BLUE}Running test suite...${NC}"
    echo "===================="

    local test_script="$PROJECT_ROOT/tests/run.sh"

    if [ -f "$test_script" ]; then
        cd "$PROJECT_ROOT/tests"
        ./run.sh $model_arg $transcribe_model_arg
    else
        echo -e "${RED}Error: Test script not found at $test_script${NC}"
        exit 1
    fi
}

# Main command handling
if [ $# -eq 0 ]; then
    print_usage
    exit 1
fi

COMMAND="$1"
MODEL_ID="${2:-$DEFAULT_MODEL_ID}"

# Check dependencies for commands that need them
case "$COMMAND" in
    build|download|run)
        check_dependencies
        ;;
esac

case "$COMMAND" in
    build)
        build_chat
        ;;
    download)
        download_model "$MODEL_ID"
        ;;
    run)
        run_chat "$MODEL_ID"
        ;;
    test)
        run_tests "$@"
        ;;
    -h|--help|help)
        print_usage
        ;;
    *)
        echo -e "${RED}Error: Unknown command '$COMMAND'${NC}"
        echo ""
        print_usage
        exit 1
        ;;
esac
